//==========================================================================
// Copyright (c) 2000-2008,  Elastos, Inc.  All Rights Reserved.
//==========================================================================
#include <aura.h>
#include "win32.h"
#include "private.h"
#include "cgmodule.h"
#include "ntimage.h"
#include "win32module.h"
#include <string.h>

void UnicToAnsi(char *pszName, const wchar_t *pwszName);
void AnsiToUnic(wchar_t *pwszName, const char *pszName);


void AnsiToUnic(wchar_t *pwszName, const char *pszName)
{
    while (*pszName) {
        *pwszName++ = (wchar_t)*pszName++;
    }
    *pwszName = 0;
}

extern DLinkNode s_ModuleList;
Aura_ECode GlobalLoadModuleA(
        CHAR *szName,
        Aura_UInt32 dwFlags,
        ModuleFileType  modFileType,
        GModule ** ppModule);

GModule * FindModuleByName(const Aura_PWChar wszName)
{
    GModule * pModule;
    ForEachDLinkNode(GModule *, pModule, &s_ModuleList) {
        if (!wcscmp((const wchar_t*)wszName, (const wchar_t*)(pModule->m_wszFullPath))) {
            pModule->AddRef();
            return pModule;
        }
    }
    return AURA_NULL;
}

GModule * FindModuleByImageBase(Aura_UInt32 uImageBase)
{
    GModule * pModule;
    ForEachDLinkNode(GModule *, pModule, &s_ModuleList) {
        if (uImageBase == pModule->m_uImageBase) {
                pModule->AddRef();
                return pModule;
        }
    }
    return AURA_NULL;
}



Aura_ECode FixupExport(
        GModule * pModule,
        Aura_UInt32 uImageBase)
{   
    PIMAGE_EXPORT_DIRECTORY pExport;
    register Aura_UInt32 n, ** ppLocal, uExportRVA;

    pExport = (PIMAGE_EXPORT_DIRECTORY)(uImageBase
                + pModule->m_pExport->m_uVAddr
                + pModule->m_uExportOffset);
                               
    uExportRVA = pModule->m_pExport->m_uVAddr + pModule->m_uExportOffset;            

    //
    // Fixup the exported functions' address.
    //
    ppLocal = (Aura_UInt32 **)(uImageBase + (Aura_UInt32)pExport->AddressOfFunctions);
    
    pExport->AddressOfFunctions = (Aura_UInt32)ppLocal;
   
    for (n = 0; n < pExport->NumberOfFunctions; n++) {
        if (0 != ppLocal[n]) {
            ppLocal[n] = (Aura_UInt32 *)((Aura_UInt32)ppLocal[n] + uImageBase);
//            Aura_printf("Export func%d fixed add: 0x%x.\n", n, ppLocal[n]);
        }
    }

    //
    // Fixup the address of name of exported functions.
    //
    ppLocal = (Aura_UInt32 **)(uImageBase + (Aura_UInt32)pExport->AddressOfNames);
    pExport->AddressOfNames = (Aura_UInt32)ppLocal;
    
   
    for (n = 0; n < pExport->NumberOfNames; n++) {
        if (0 != ppLocal[n]) {
            ppLocal[n] = (Aura_UInt32 *)((Aura_UInt32)ppLocal[n] + uImageBase);
//            Aura_printf("Export func%d name fixed add: 0x%x.\n", n, ppLocal[n]);
//            Aura_printf("Export func%d name : %s.\n", n, (CHAR *)ppLocal[n]);
        }
    }

    //
    // The ordinals base address is fixuped up
    //
    pExport->AddressOfNameOrdinals
        = (Aura_UInt32)((Aura_UInt32)pExport->AddressOfNameOrdinals + uImageBase);

//    Aura_printf("Export name ordinals fixed add: 0x%x.\n", pExport->AddressOfNameOrdinals);
    
    return AURA_NOERROR;
}


Aura_UInt32 GetExportAddressByOrdinal(
        GModule * pModule,
        Aura_UInt32 uOrdinal)
{
    PIMAGE_EXPORT_DIRECTORY pExport;

//    assert(ModuleFormatType_PE == pModule->m_formatType);

    if (AURA_NULL == pModule->m_pExport) return 0;

    pExport = (PIMAGE_EXPORT_DIRECTORY)(
                pModule->m_uImageBase
                + pModule->m_pExport->m_uVAddr
                + pModule->m_uExportOffset);


    uOrdinal -= pExport->Base;

    if (uOrdinal < pExport->NumberOfFunctions) {
        Aura_PUInt32 pAddress  = (Aura_PUInt32)(pExport->AddressOfFunctions);
        return (Aura_UInt32)pAddress[uOrdinal];
    }
    else {
        return 0;
    }
}

Aura_UInt32 GetExportAddressByName(
        GModule * pModule,
        const CHAR * pszName)
{
    PIMAGE_EXPORT_DIRECTORY pExport;
    Aura_UInt32 n;
    CHAR ** ppNames;

    if (AURA_NULL == pModule->m_pExport) return 0;

    pExport = (PIMAGE_EXPORT_DIRECTORY)(
                pModule->m_uImageBase
                + pModule->m_pExport->m_uVAddr
                + pModule->m_uExportOffset);

    ppNames = (CHAR **)pExport->AddressOfNames;    
   
    for (n = 0; n < pExport->NumberOfNames; n++) {
        if (!strcmp(pszName, ppNames[n])) {
            Aura_PUInt16 pAddressOrdinals = (Aura_PUInt16)(pExport->AddressOfNameOrdinals);
            
            n = pAddressOrdinals[n];
            
            Aura_PUInt32 pAddressFun = (Aura_PUInt32)(pExport->AddressOfFunctions);
            
            return (Aura_UInt32)(pAddressFun[n]);
        }
    }
    return 0;
}



Aura_ECode Win32Module_LoadA(
        /* [in] */ Aura_PChar szName,
        /* [in] */ Aura_ModuleLoadFlags flags,
        /* [out] */ Aura_Handle * phModule)
{
    Aura_WChar wszName[256] = {0};

    AnsiToUnic((wchar_t *)wszName, szName);
 
    HINSTANCE hModule = NULL;

    hModule = LoadLibraryW((LPCWSTR)(const wchar_t *)wszName);
    if (NULL == hModule) {
        goto E_FAIL_EXIT;
    }

    *phModule = hModule;
    return AURA_NOERROR;

E_FAIL_EXIT:
    return GetLastECodeError();
}


Aura_ECode LoadOneImportFromExtern(GModule *pModule, PIMAGE_IMPORT_DESCRIPTOR pImport, Aura_UInt32 uImageBase) 
{
    Aura_UInt32 uFunction;
    CHAR * pszDllName = (CHAR *)(uImageBase + pImport->Name);
    
    Aura_PVoid hDll;
    
//    Aura_printf("pszDllName = %s \n", pszDllName);
    Aura_ECode ec = Win32Module_LoadA(pszDllName, 0, &hDll);
    if(AURA_FAILED(ec)) {
        Aura_printf("AuraModule_LoadA error, ec = %d. FILE:%s,LINE:%d\n", 
        		ec, __FILE__, __LINE__);
        return ec;
    }
    
    Win32ModuleRef *pWin32ModuleRef = new Win32ModuleRef(hDll);
    if (!pWin32ModuleRef) {
    	Win32Module_Unload(hDll);
        return AURA_E_OUT_OF_MEMORY;
    }
    
    pModule->m_win32ModuleList.InsertPrev(pWin32ModuleRef);
   
    PIMAGE_THUNK_DATA pOrgThunk, pFirstThunk;
    pOrgThunk = (PIMAGE_THUNK_DATA)(uImageBase + pImport->OriginalFirstThunk);
    pFirstThunk = (PIMAGE_THUNK_DATA)(uImageBase + (Aura_UInt32)pImport->FirstThunk);
    
    while (*(Aura_PUInt32)pOrgThunk != 0) {
        if (IMAGE_SNAP_BY_ORDINAL(pOrgThunk->u1.Ordinal)) {
            ec = Win32Module_GetProcAddressByOrdinal(hDll, \
                    (Aura_Int32)(pOrgThunk->u1.Ordinal & 0x0000ffff), &uFunction);
            if(AURA_FAILED(ec)) {
                return ec;
            }                 
        } else {
            PIMAGE_IMPORT_BY_NAME pImportName = (PIMAGE_IMPORT_BY_NAME)( 
            (Aura_UInt32)pOrgThunk->u1.AddressOfData + uImageBase); 
            ec = Win32Module_GetProcAddressByName(hDll, \
                    (Aura_PChar)(pImportName->Name), &uFunction);
            if(AURA_FAILED(ec)) {
                return ec;
            } 
        }
        
        if (!uFunction) {
            return AURA_E_DOES_NOT_EXIST;
        }
        
        pFirstThunk->u1.Function = (Aura_UInt32)uFunction;

        pOrgThunk++;
        pFirstThunk++;
    }
    
    return AURA_NOERROR;
}


Aura_ECode LoadOneImportFromIntern(GModule *pModule, PIMAGE_IMPORT_DESCRIPTOR pImport, Aura_UInt32 uImageBase)
{
    GModule *pExpModule;
    Aura_UInt32 uFunction;
    
    CHAR * pszDllName = (CHAR *)(uImageBase + pImport->Name);
    
    Aura_ECode ec = GlobalLoadModuleA(pszDllName, 0, ModuleFileType_DLL, &pExpModule);
    if (AURA_NOERROR != ec) {
//        Aura_printf("GlobalLoadModuleA error, ec = %x. FILE: %s, LINE:%d\n", 
//                ec, __FILE__, __LINE__);
        return ec;
    }
    
    ModuleRef * pModuleRef = new ModuleRef(pExpModule);
    if (!pModuleRef) {
        return AURA_E_OUT_OF_MEMORY;
    }
    
    pModule->m_rModuleList.InsertPrev(pModuleRef);
    
    PIMAGE_THUNK_DATA pOrgThunk, pFirstThunk;
    
    pOrgThunk = (PIMAGE_THUNK_DATA)(uImageBase + pImport->OriginalFirstThunk);
    pFirstThunk = (PIMAGE_THUNK_DATA)(uImageBase + (Aura_UInt32)pImport->FirstThunk);
    
    while (*(Aura_PUInt32)pOrgThunk != 0) {
        if (IMAGE_SNAP_BY_ORDINAL(pOrgThunk->u1.Ordinal)) {
            uFunction = GetExportAddressByOrdinal(
                    pExpModule,
                    IMAGE_ORDINAL(pOrgThunk->u1.Ordinal));
            
        } else {
             PIMAGE_IMPORT_BY_NAME pImportName = (PIMAGE_IMPORT_BY_NAME)( \
             (Aura_UInt32)pOrgThunk->u1.AddressOfData + uImageBase);
             uFunction = GetExportAddressByName(
                    pExpModule,
                    (CHAR *)pImportName->Name);
        }
        
        if (!uFunction) {
            return AURA_E_DOES_NOT_EXIST;
        }
        
        pFirstThunk->u1.Function = (Aura_UInt32)uFunction;

        pOrgThunk++;
        pFirstThunk++;
    }
    
    return AURA_NOERROR;   
}


Aura_ECode FixupImport(
        GModule *pModule,
        Aura_UInt32 uImageBase)
{
    PIMAGE_IMPORT_DESCRIPTOR pImport; 
    Aura_ECode ec;
    
    pImport = (PIMAGE_IMPORT_DESCRIPTOR)
                    (uImageBase
                    + pModule->m_pImport->m_uVAddr
                    + pModule->m_uImportOffset);
                    
    
     while (pImport->OriginalFirstThunk != 0) {
        //
        // Get the dependent DLL's name and try to load it.
        //
        ec = LoadOneImportFromIntern(pModule, pImport, uImageBase);
        if(AURA_FAILED(ec)) {
            if(ec == AURA_E_FILE_NOT_FOUND) {
                ec = LoadOneImportFromExtern(pModule, pImport, uImageBase);
                if(AURA_FAILED(ec)) {
                    Aura_printf("LoadOneImportFromSys error.ec = %x.  FILE:%s, LINE:%d\n", 
                    		ec, __FILE__, __LINE__);
                    return ec;
                }
            } else {
                Aura_printf("LoadOneImportFromNotSys error.ec =%x. FILE:%s, LINE:%d\n", 
                		ec, __FILE__, __LINE__);
                return ec;
            }
        }
        
        pImport++;
    }

    return AURA_NOERROR;                
}   


Aura_ECode RelocPeModule(
        GModule * pModule,
        Aura_UInt32 uOffset,
        Aura_UInt32 uRelocOffset)
{
    DWORD TpOffset;
    PDWORD lpLink;
    
    PIMAGE_BASE_RELOCATION pReloc = (PIMAGE_BASE_RELOCATION) ((DWORD)uRelocOffset +  \
    		(DWORD)pModule->m_uImageBase);
    
    while(pReloc->VirtualAddress != 0) {
        PWORD lpTypeOffset = (PWORD) ((DWORD)pReloc + sizeof(IMAGE_BASE_RELOCATION));
		while (lpTypeOffset < (PWORD)((DWORD)pReloc + (DWORD)pReloc->SizeOfBlock)) {
			TpOffset = *lpTypeOffset & 0xF000;
			
			if(TpOffset == 0x3000) {
				lpLink = (PDWORD) ((DWORD)pModule->m_uImageBase + pReloc->VirtualAddress + 
				(*lpTypeOffset & 0xFFF));
				*lpLink = (DWORD)uOffset + (*lpLink);
			}
			else {
				if (TpOffset != 0) {
					return AURA_E_FAIL;
				}
			}
			lpTypeOffset++;
		}
		
		pReloc = (PIMAGE_BASE_RELOCATION) ((DWORD)pReloc + (DWORD)pReloc->SizeOfBlock);
    }
    
    return AURA_NOERROR;
}